# utils/dataset_validator.py

import os
import yaml
import shutil
from pathlib import Path

"""
    数据集验证工具类
    包含以下功能：
    - 验证数据集目录结构
    - 验证YAML配置文件    
    - 验证图像与标注文件对应关系
"""
class DatasetValidator:
    @staticmethod
    def validate_structure(dataset_path):
        """
        验证数据集目录结构
        """
        required_dirs = [
            'images/train',
            'images/val',
            'labels/train',
            'labels/val'
        ]

        missing_dirs = []
        dataset_path = Path(dataset_path)

        for rel_path in required_dirs:
            parts = rel_path.split('/')
            current_dir = dataset_path
            valid = True

            for part in parts:
                found = False
                for item in current_dir.iterdir():
                    if item.name.startswith(('.', '__')):
                        continue
                    if item.is_dir() and item.name.lower() == part:
                        current_dir = item
                        found = True
                        break
                if not found:
                    valid = False
                    break

            if not valid:
                missing_dirs.append(rel_path)

        if missing_dirs:
            raise ValueError(f'缺失必要目录: {", ".join(missing_dirs)}')

    @staticmethod
    def validate_yaml(dataset_path):
        """
        验证YAML配置文件并返回(yaml路径, 类别数)
        """
        yaml_files = [f for f in Path(dataset_path).glob('*.yaml')] + \
                     [f for f in Path(dataset_path).glob('*.yml')]

        if not yaml_files:
            raise ValueError('未找到YAML配置文件')
        if len(yaml_files) > 1:
            raise ValueError('发现多个YAML文件')

        yaml_path = yaml_files[0]
        try:
            with open(yaml_path, 'r') as f:
                content = yaml.safe_load(f)
                required_keys = ['train', 'val', 'nc', 'names']

                # 检查必需字段
                missing_keys = [k for k in required_keys if k not in content]
                if missing_keys:
                    raise ValueError(f'YAML文件缺少字段: {", ".join(missing_keys)}')

                # 验证类别数量一致性
                nc = content.get('nc', 0)
                names = content.get('names', [])
                if nc != len(names):
                    raise ValueError(f'类别数不一致: nc={nc}, names数量={len(names)}')

                return str(yaml_path), nc
        except yaml.YAMLError as e:
            raise ValueError(f'YAML解析错误: {str(e)}')

    @staticmethod
    def validate_pairs(dataset_path):
        """验证图像与标注对应关系并统计有效样本"""
        def get_valid_pairs(img_dir, label_dir):
            # 获取所有图片和标签的基名（不带扩展名）
            image_files = {
                f.stem.lower(): f
                for f in img_dir.glob('*')
                if f.suffix.lower() in ['.jpg', '.jpeg', '.png']
            }

            label_files = {
                f.stem.lower(): f
                for f in label_dir.glob('*.txt')
            }

            # 找出共同基名的文件
            common_stems = set(image_files.keys()) & set(label_files.keys())
            return len(common_stems)

        try:
            train_img_dir = Path(dataset_path) / 'images' / 'train'
            train_label_dir = Path(dataset_path) / 'labels' / 'train'
            val_img_dir = Path(dataset_path) / 'images' / 'val'
            val_label_dir = Path(dataset_path) / 'labels' / 'val'

            return {
                'train_samples': get_valid_pairs(train_img_dir, train_label_dir),
                'val_samples': get_valid_pairs(val_img_dir, val_label_dir)
            }
        except Exception as e:
            raise ValueError(f'文件验证失败: {str(e)}')